using System;
using System.Collections.Generic;
using System.IO;

namespace PeDALS.Tools.AZ.PSTArchive.PSTData
{
    public class TreeNodeEntry : TreeEntry
    {
        private UInt64 m64NIDraw;
        private UInt64 m64BIDrawData;
        private long mlngDataAddress;
        private short mshtDataLength;
        private UInt64 m64BIDrawSub;
        private long mlngSubAddress;
        private short mshtSubLength;
        private UInt64 m64NIDrawParent;
        private List<string> mlstrPath = new List<string>();
        private SortedList<UInt32, byte[]> mslbaSubNodes = new SortedList<uint, byte[]>();
        private SortedList<UInt32, PSTHeap> mslphHeaps = new SortedList<uint, PSTHeap>();

        public TreeNodeEntry(UInt64 ui64NIDraw, UInt64 ui64BIDrawData, UInt64 ui64NIDrawSub, UInt64 ui64NIDrawParent, string strPath, long lngEntryAddress)
            : base(strPath, lngEntryAddress)
        {
            m64NIDraw = ui64NIDraw;
            m64BIDrawData = ui64BIDrawData;
            m64BIDrawSub = ui64NIDrawSub;
            m64NIDrawParent = ui64NIDrawParent;
            mlstrPath.Add(this.EntryPath);
        }

        public UInt64 NIDraw
        {
            get { return m64NIDraw; }
        }

        public UInt64 NID
        {
            get { return m64NIDraw >> 5; }
        }

        public NIDTypes NIDType
        {
            get { return (NIDTypes) (m64NIDraw & 0x1f); }
        }

        public UInt64 BIDrawData
        {
            get { return m64BIDrawData; }
        }

        public long DataAddress
        {
            get { return mlngDataAddress; }
        }

        public short DataLength
        {
            get { return mshtDataLength; }
        }

        public UInt64 BIDrawSub
        {
            get { return m64BIDrawSub; }
        }

        public long SubAddress
        {
            get { return mlngSubAddress; }
        }

        public short SubLength
        {
            get { return mshtSubLength; }
        }

        public UInt64 NIDrawParent
        {
            get { return m64NIDrawParent; }
        }

        public UInt64 NIDParent
        {
            get { return m64NIDrawParent >> 5; }
        }

        public NIDTypes NIDParentType
        {
            get { return (NIDTypes) (m64NIDrawParent & 0x1f); }
        }

        public List<string> PathList
        {
            get { return mlstrPath; }
        }

        internal SortedList<UInt32, byte[]> ByteData
        {
            get { return mslbaSubNodes; }
            set { mslbaSubNodes = value; }
        }

        internal SortedList<UInt32, PSTHeap> Heaps 
        {
            get { return mslphHeaps; }
            set { mslphHeaps = value; }
        }

        public long LoadData(BinaryReader brReader, bool blnUnicode, SortedList<UInt64, TreeBlockEntry> sltbeBlockEntries, byte bCryptMethod) 
        {
            SortedList<UInt32, List<byte[]>> sllbaBlocks = new SortedList<uint, List<byte[]>>();

            if (sltbeBlockEntries.ContainsKey(m64BIDrawData))
            {
                TreeBlockEntry tbeEntry = sltbeBlockEntries[m64BIDrawData];
                mlngDataAddress = tbeEntry.BlockAddress;
                mshtDataLength = tbeEntry.BlockLength;
                sllbaBlocks = LoadBlocks(brReader, blnUnicode, tbeEntry, sltbeBlockEntries, (UInt32) NIDType, bCryptMethod);
            }
            if (sltbeBlockEntries.ContainsKey(m64BIDrawSub))
            {
                TreeBlockEntry tbeSubEntry = sltbeBlockEntries[m64BIDrawSub];
                mlngSubAddress = tbeSubEntry.BlockAddress;
                mshtSubLength = tbeSubEntry.BlockLength;
                SortedList<UInt32, List<byte[]>> sllbaSubBlocks = LoadBlocks(brReader, blnUnicode, tbeSubEntry, sltbeBlockEntries, (UInt32)NIDType, bCryptMethod);

                foreach (KeyValuePair<UInt32, List<byte[]>> kvpBlockEntry in sllbaSubBlocks)
                {
                    if (sllbaBlocks.ContainsKey(kvpBlockEntry.Key))
                    {
                        sllbaBlocks[kvpBlockEntry.Key].AddRange(kvpBlockEntry.Value);
                    }
                    else
                    {
                        sllbaBlocks.Add(kvpBlockEntry.Key, kvpBlockEntry.Value);
                    }
                }
            }
            foreach (KeyValuePair<UInt32, List<byte[]>> kvpBlockEntry in sllbaBlocks)
            {
                if (kvpBlockEntry.Value[0][2] == 0xEC)
                {
                    PSTHeap phHeap = new PSTHeap(kvpBlockEntry.Value);
                    if (!mslphHeaps.ContainsKey(kvpBlockEntry.Key))
                    {
                        mslphHeaps.Add(kvpBlockEntry.Key, new PSTHeap(kvpBlockEntry.Value));
                    }
                    else
                    {
                        System.Diagnostics.Trace.Write(string.Format("*******NID 0x{0:x8} ERROR:  Heap Exists for HNID 0x{1:x8}", NIDraw, kvpBlockEntry.Key));
                    }
                }
                else
                {
                    foreach (byte[] baBlock in kvpBlockEntry.Value)
                    {
                        if (mslbaSubNodes.ContainsKey(kvpBlockEntry.Key))
                        {
                            byte[] baEntry = mslbaSubNodes[kvpBlockEntry.Key];
                            int intOffset = baEntry.Length;
                            Array.Resize(ref baEntry, baEntry.Length + baBlock.Length);
                            Array.Copy(baBlock, 0, baEntry, intOffset, baBlock.Length);
                            mslbaSubNodes[kvpBlockEntry.Key] = baEntry;
                        }
                        else
                        {
                            mslbaSubNodes.Add(kvpBlockEntry.Key, baBlock);
                        }
                    }
                }
            }
            sllbaBlocks = null;
            foreach(KeyValuePair<UInt32, PSTHeap> kvpHeap in mslphHeaps)
            {
                kvpHeap.Value.ProcessHeap(NIDraw, kvpHeap.Key, mslbaSubNodes, blnUnicode);
            }
            if (mslphHeaps.Count > 0) mslbaSubNodes = null;

            return GC.GetTotalMemory(false);
        }

        private SortedList<UInt32, List<byte[]>> LoadBlocks(BinaryReader brReader, bool blnUnicode, TreeBlockEntry tbeEntry, SortedList<UInt64, TreeBlockEntry> sltbeEntries, UInt32 u32HNID, byte bCryptMethod)
        {
            byte bType = 0;
            byte bLevel = 0;
            short shtSubBlockEntries = 0;
            int intReferenceLength = -1;
            long lngCurrentAddress = 0;
            long lngTrailerAddress = 0;
            SortedList<UInt32, List<byte[]>> slbaResults = new SortedList<UInt32, List<byte[]>>();
            bool blnValidBlock = true;

            int intRemainder = (tbeEntry.BlockLength + (blnUnicode ? 16 : 12)) % 64;
            int intBlockPadding = intRemainder > 0 ? 64 - intRemainder : intRemainder;

            brReader.BaseStream.Seek(tbeEntry.BlockAddress + tbeEntry.BlockLength, SeekOrigin.Begin);
            byte[] baPadding = brReader.ReadBytes(intBlockPadding);
            lngTrailerAddress = brReader.BaseStream.Position;
            PSTBlockTrailer btTrailer = new PSTBlockTrailer(brReader, blnUnicode);
            if (btTrailer.BIDraw != tbeEntry.BIDraw)
            {
                blnValidBlock = false;
                System.Diagnostics.Trace.Write(string.Format("*******NID 0x{0:x8} ERROR:  *********\nBIDraw:  0x{1:x8}; Address:  0x{2:x8}; Length:  0x{3:x8}; Entry Address:  0x{4:x8}; Path:  {5}\nTrailer:  BIDraw:  0x{6:x8}; Address:  0x{7:x8}; Remainder:  {8}; Padding:  {9}", this.NIDraw, tbeEntry.BIDraw, tbeEntry.BlockAddress, tbeEntry.BlockLength, tbeEntry.EntryAddress, tbeEntry.EntryPath, btTrailer.BIDraw, lngTrailerAddress, intRemainder, intBlockPadding));
                List<TreeBlockEntry> ltbeEntries = new List<TreeBlockEntry>();
                ltbeEntries.AddRange(sltbeEntries.Values);
                List<TreeBlockEntry> ltbeAddressMatches = ltbeEntries.FindAll(delegate(TreeBlockEntry tbeFoundEntry) { return tbeFoundEntry.BlockAddress == tbeEntry.BlockAddress & tbeFoundEntry.BIDraw != tbeEntry.BIDraw; });
                foreach (TreeBlockEntry tbeAddressMatch in ltbeAddressMatches)
                {
                    System.Diagnostics.Trace.Write(string.Format("Block Address Match:  BIDraw:  0x{0:x8}; Address:  0x{1:x8}; Length:  0x{2:x8}; Entry Address:  0x{3:x8}; Path:  {4}", tbeAddressMatch.BIDraw, tbeAddressMatch.BlockAddress, tbeAddressMatch.BlockLength, tbeAddressMatch.EntryAddress, tbeAddressMatch.EntryPath));
                }
            }
            brReader.BaseStream.Seek(tbeEntry.BlockAddress, SeekOrigin.Begin);
            if (tbeEntry.Internal)
            {
                bType = brReader.ReadByte();
                bLevel = brReader.ReadByte();
                shtSubBlockEntries = brReader.ReadInt16();
                List<UInt64> lu64BIDraw = new List<UInt64>();
                UInt64 u64BIDraw;
                if (bType == 0x01)
                {
                    intReferenceLength = brReader.ReadInt32();
                    do
                    {
                        u64BIDraw = blnUnicode ? brReader.ReadUInt64() : Convert.ToUInt64(brReader.ReadUInt32());
                        if (u64BIDraw > 0) lu64BIDraw.Add(u64BIDraw);
                    }
                    while (lu64BIDraw.Count < shtSubBlockEntries);
                    lngCurrentAddress = brReader.BaseStream.Position;
                    foreach (UInt64 u64LoadBIDraw in lu64BIDraw)
                    {
                        if (sltbeEntries.ContainsKey(u64LoadBIDraw & 0xFFFFFFFE))
                        {
                            TreeBlockEntry tbeLoadEntry = sltbeEntries[u64LoadBIDraw & 0xFFFFFFFE];
                            SortedList<UInt32, List<byte[]>> slbaEntry = LoadBlocks(brReader, blnUnicode, tbeLoadEntry, sltbeEntries, u32HNID, bCryptMethod);
                            foreach (KeyValuePair<UInt32, List<byte[]>> kvpEntry in slbaEntry)
                            {
                                if (slbaResults.ContainsKey(kvpEntry.Key))
                                {
                                    foreach (byte[] baValue in (List<byte[]>)kvpEntry.Value)
                                    {
                                        ((List<byte[]>)slbaResults[kvpEntry.Key]).Add(baValue);
                                    }
                                }
                                else
                                {
                                    slbaResults.Add(kvpEntry.Key, kvpEntry.Value);
                                }
                            }
                        }
                        else
                        {
                            if (u64LoadBIDraw != 0)
                            {
                                System.Diagnostics.Trace.Write(string.Format("*******NID 0x{0:x8} Block load failed:  Block not found BIDraw:  0x{1:x8}; Reserved Bit Masked BID Entry:  0x{2:x8}; {3} Node", NIDraw, u64LoadBIDraw, u64LoadBIDraw & 0xFFFFFFFE, (u64LoadBIDraw & 0x2) == 0x2 ? "Internal" : "External"));
                            }
                        }
                    }
                    brReader.BaseStream.Seek(lngCurrentAddress, SeekOrigin.Begin);
                }
                else
                {
                    if (blnUnicode) intReferenceLength = brReader.ReadInt32();
                    int intPadding;
                    if (bLevel == 0x00)
                    {
                        List<SLEntry> lsleEntries = new List<SLEntry>();
                        UInt64 u64slNIDraw;
                        UInt64 u64slBIDrawData;
                        UInt64 u64slBIDrawSub;
                        SLEntry sleEntry;
                        do
                        {
                            u64slNIDraw = Convert.ToUInt64(brReader.ReadUInt32());
                            if (blnUnicode) intPadding = brReader.ReadInt32();
                            u64slBIDrawData = blnUnicode ? brReader.ReadUInt64() : Convert.ToUInt64(brReader.ReadUInt32());
                            u64slBIDrawSub = blnUnicode ? brReader.ReadUInt64() : Convert.ToUInt64(brReader.ReadUInt32());

                            sleEntry = new SLEntry(u64slNIDraw, u64slBIDrawData, u64slBIDrawSub);
                            if (sleEntry.NIDraw > 0) lsleEntries.Add(sleEntry);
                        }
                        while (lsleEntries.Count < shtSubBlockEntries);
                        lngCurrentAddress = brReader.BaseStream.Position;
                        foreach (SLEntry sleLoadEntry in lsleEntries)
                        {
                            if (sltbeEntries.ContainsKey(sleLoadEntry.BIDrawData & 0xFFFFFFFE))
                            {
                                TreeBlockEntry tbeLoadsleEntry = sltbeEntries[sleLoadEntry.BIDrawData & 0xFFFFFFFE];
                                SortedList<UInt32, List<byte[]>> slbaEntry = LoadBlocks(brReader, blnUnicode, tbeLoadsleEntry, sltbeEntries, Convert.ToUInt32(sleLoadEntry.NIDraw), bCryptMethod);
                                foreach (KeyValuePair<UInt32, List<byte[]>> kvpEntry in slbaEntry)
                                {
                                    if (slbaResults.ContainsKey(kvpEntry.Key))
                                    {
                                        foreach (byte[] baValue in (List<byte[]>)kvpEntry.Value)
                                        {
                                            ((List<byte[]>)slbaResults[kvpEntry.Key]).Add(baValue);
                                        }
                                    }
                                    else
                                    {
                                        slbaResults.Add(kvpEntry.Key, kvpEntry.Value);
                                    }
                                }
                            }
                            else
                            {
                                if (sleLoadEntry.BIDrawData != 0)
                                {
                                    System.Diagnostics.Trace.Write(string.Format("*******NID 0x{0:x8} Block load failed:  Block not found BIDraw:  0x{1:x8}; Reserved Bit Masked BID Entry:  0x{2:x8}; {3} Node", NIDraw, sleLoadEntry.BIDrawData, sleLoadEntry.BIDrawData & 0xFFFFFFFE, (sleLoadEntry.BIDrawData & 0x2) == 0x2 ? "Internal" : "External"));
                                }
                            }
                            if (sltbeEntries.ContainsKey(sleLoadEntry.BIDrawSub & 0xFFFFFFFE))
                            {
                                TreeBlockEntry tbeLoadsleSubEntry = sltbeEntries[sleLoadEntry.BIDrawSub & 0xFFFFFFFE];
                                SortedList<UInt32, List<byte[]>> slbaEntry = LoadBlocks(brReader, blnUnicode, tbeLoadsleSubEntry, sltbeEntries, Convert.ToUInt32(sleLoadEntry.NIDraw), bCryptMethod);
                                foreach (KeyValuePair<UInt32, List<byte[]>> kvpEntry in slbaEntry)
                                {
                                    if (slbaResults.ContainsKey(kvpEntry.Key))
                                    {
                                        foreach (byte[] baValue in (List<byte[]>)kvpEntry.Value)
                                        {
                                            ((List<byte[]>)slbaResults[kvpEntry.Key]).Add(baValue);
                                        }
                                    }
                                    else
                                    {
                                        slbaResults.Add(kvpEntry.Key, kvpEntry.Value);
                                    }
                                }
                            }
                            else
                            {
                                if (sleLoadEntry.BIDrawSub != 0)
                                {
                                    System.Diagnostics.Trace.Write(string.Format("*******NID 0x{0:x8} Block load failed:  Block not found BIDraw:  0x{1:x8}; Reserved Bit Masked BID Entry:  0x{2:x8}; {3} Node", NIDraw, sleLoadEntry.BIDrawSub, sleLoadEntry.BIDrawSub & 0xFFFFFFFE, (sleLoadEntry.BIDrawSub & 0x2) == 0x2 ? "Internal" : "External"));
                                }
                            }
                        }
                        brReader.BaseStream.Seek(lngCurrentAddress, SeekOrigin.Begin);
                    }
                    else
                    {
                        List<SIEntry> lsieEntries = new List<SIEntry>();
                        UInt64 u64siNIDraw;
                        UInt64 u64siBIDraw;
                        SIEntry sieEntry;
                        do
                        {
                            u64siNIDraw = Convert.ToUInt64(brReader.ReadInt32());
                            if (blnUnicode) intPadding = brReader.ReadInt32();
                            u64siBIDraw = blnUnicode ? brReader.ReadUInt64() : Convert.ToUInt64(brReader.ReadUInt32());
                            sieEntry = new SIEntry(u64siNIDraw, u64siBIDraw);
                            if (u64siNIDraw > 0) lsieEntries.Add(sieEntry);
                        }
                        while (lsieEntries.Count < shtSubBlockEntries);
                        lngCurrentAddress = brReader.BaseStream.Position;
                        foreach (SIEntry sieLoadEntry in lsieEntries)
                        {
                            if (sltbeEntries.ContainsKey(sieLoadEntry.BIDraw & 0xFFFFFFFE))
                            {
                                TreeBlockEntry tbeLoadsieEntry = sltbeEntries[sieLoadEntry.BIDraw & 0xFFFFFFFE];
                                SortedList<UInt32, List<byte[]>> slbaEntry = LoadBlocks(brReader, blnUnicode, tbeLoadsieEntry, sltbeEntries, Convert.ToUInt32(sieLoadEntry.NIDraw), bCryptMethod);
                                foreach (KeyValuePair<UInt32, List<byte[]>> kvpEntry in slbaEntry)
                                {
                                    if (slbaResults.ContainsKey(kvpEntry.Key))
                                    {
                                        foreach (byte[] baValue in (List<byte[]>)kvpEntry.Value)
                                        {
                                            ((List<byte[]>)slbaResults[kvpEntry.Key]).Add(baValue);
                                        }
                                    }
                                    else
                                    {
                                        slbaResults.Add(kvpEntry.Key, kvpEntry.Value);
                                    }
                                }
                            }
                            else
                            {
                                if (sieLoadEntry.BIDraw != 0)
                                {
                                    System.Diagnostics.Trace.Write(string.Format("*******NID 0x{0:x8} Block load failed:  Block not found BIDraw:  0x{1:x8}; Reserved Bit Masked BID Entry:  0x{2:x8}; {3} Node", NIDraw, sieLoadEntry.BIDraw, sieLoadEntry.BIDraw & 0xFFFFFFFE, (sieLoadEntry.BIDraw & 0x2) == 0x2 ? "Internal" : "External"));
                                }
                            }
                        }
                        brReader.BaseStream.Seek(lngCurrentAddress, SeekOrigin.Begin);
                    }
                }
                if (!blnValidBlock)
                {
                    System.Diagnostics.Trace.Write(string.Format("Type:  0x{0:X2}; Level:  0x{1:X2}; SubBlockEntries:  {2}; Reference Length:  0x{3:x8}({3})", bType, bLevel, shtSubBlockEntries, intReferenceLength));
                }
            }
            else
            {
                List<byte[]> lbaNewList = new List<byte[]>();
                lbaNewList.Add(DecryptBlock(brReader.ReadBytes(tbeEntry.BlockLength), 0, bCryptMethod));
                slbaResults.Add(u32HNID, lbaNewList);
            }
            return slbaResults;
        }

        private byte[] DecryptBlock(byte[] baBlock, int intBID, byte bAlgorithm)
        {
            byte[] baResults = (baBlock != null) ? Crypto.Decrypt(baBlock, baBlock.Length, intBID, bAlgorithm) : null;
            return baResults;
        }
    }
}